import torch.nn

from scaler import Scaler


class Lvminn(torch.nn.Module):
    def __init__(self,
                 t_scaler: Scaler = Scaler(-1, 1),
                 hidden_size: int = 300, hidden_count: int = 3):
        super().__init__()

        self.network: torch.nn.Sequential = torch.nn.Sequential(
            t_scaler,
            torch.nn.Linear(1, hidden_size),
            torch.nn.Tanh(),
            *[
                 torch.nn.Linear(hidden_size, hidden_size),
                 torch.nn.Tanh()
             ] * hidden_count,
            torch.nn.Linear(hidden_size, 2)
        )

        self.a1: torch.nn.Parameter = torch.nn.Parameter(torch.tensor(0.))
        self.b1: torch.nn.Parameter = torch.nn.Parameter(torch.tensor(0.))
        self.a2: torch.nn.Parameter = torch.nn.Parameter(torch.tensor(0.))
        self.b2: torch.nn.Parameter = torch.nn.Parameter(torch.tensor(0.))

    def forward(self, t_batch: torch.Tensor) -> torch.Tensor:
        x = torch.stack([t_batch], dim=-1)
        return self.network(x)
